#!/usr/bin/env python3

import argparse
import os
import sys
import time
from threading import Thread, Event
import grpc

sys.path.append(
    os.path.join(os.path.dirname(os.path.abspath(__file__)),
                 '../../utils/'))
import p4runtime_lib.bmv2
import p4runtime_lib.helper
from p4runtime_lib.switch import ShutdownAllSwitchConnections

class ECNMonitor:
    def __init__(self, switch, p4info_helper):
        self.switch = switch
        self.p4info_helper = p4info_helper
        self._stop_event = Event()
        self.monitor_thread = None
        self.congestion_status = {}  # 初始化状态字典
        self.last_print_time = 0

    def start(self):
        """启动监控线程"""
        if not self.monitor_thread:
            self._stop_event.clear()
            self.monitor_thread = Thread(target=self._monitor)
            self.monitor_thread.daemon = True
            self.monitor_thread.start()

    def stop(self):
        """停止监控"""
        self._stop_event.set()
        if self.monitor_thread:
            self.monitor_thread.join(timeout=1)

    def _monitor(self):
        """监控拥塞状态"""
        try:
            while not self._stop_event.is_set():
                self._update_congestion_status()
                
                # 每秒打印一次状态
                current_time = time.time()
                if current_time - self.last_print_time >= 1.0:
                    self._print_status()
                    self.last_print_time = current_time
                
                time.sleep(0.5)  # 每0.5秒检查一次
        except Exception as e:
            print(f"\n监控错误: {e}")

    def _update_congestion_status(self):
        """更新拥塞状态"""
        try:
            # 模拟读取队列深度 - 实际应根据P4程序调整
            for port in [1, 2, 3]:  # 假设有3个端口
                # 这里应该是读取寄存器或计数器的实际代码
                # 为了示例，我们使用随机数据
                qdepth = int(time.time() * 1000) % 30  # 模拟队列深度变化
                
                status = 'normal'
                if qdepth >= 20:
                    status = 'Severe congestion'
                elif qdepth >= 10:
                    status = 'Moderate congestion'
                
                self.congestion_status[port] = {
                    'status': status,
                    'qdepth': qdepth
                }
                
        except Exception as e:
            print(f"\n更新状态错误: {e}")

    def _print_status(self):
        """打印当前状态"""
        sys.stdout.write("\r")  # 回到行首
        for port, data in sorted(self.congestion_status.items()):
            status = data['status']
            qdepth = data['qdepth']
            
            color_code = ""
            if status == 'Severe congestion':
                color_code = "\033[1;31m"  # 红色
            elif status == 'Moderate congestion':
                color_code = "\033[1;33m"  # 黄色
            
            sys.stdout.write(f"{color_code}port {port}: {status}({qdepth}) \033[0m")
        sys.stdout.flush()

def main(p4info_file_path, bmv2_file_path):
    # 确保日志目录存在
    if not os.path.exists('logs'):
        os.makedirs('logs', mode=0o777)
    
    p4info_helper = p4runtime_lib.helper.P4InfoHelper(p4info_file_path)
    
    monitor = None
    switch = None
    try:
        print("初始化交换机连接...")
        # 创建交换机连接
        switch = p4runtime_lib.bmv2.Bmv2SwitchConnection(
            name='s1',
            address='127.0.0.1:50051',
            device_id=0,
            proto_dump_file='logs/s1-p4runtime-requests.txt'
        )

        # 主控制器仲裁和管道配置
        switch.MasterArbitrationUpdate()
        switch.SetForwardingPipelineConfig(
            p4info=p4info_helper.p4info,
            bmv2_json_file_path=bmv2_file_path
        )

        # 启动监控
        monitor = ECNMonitor(switch, p4info_helper)
        monitor.start()

        print("\nStart monitoring network congestion status...")
        print("real-time status:")
        while True:
            time.sleep(1)
            
    except KeyboardInterrupt:
        print("\n\nstop...")
    except Exception as e:
        print(f"\nerror: {str(e)}")
    finally:
        if monitor:
            monitor.stop()
        if switch:
            ShutdownAllSwitchConnections()
        print("\nstop")

if __name__ == '__main__':
    parser = argparse.ArgumentParser(description='实时ECN拥塞监控控制器')
    parser.add_argument('--p4info', help='p4info文件路径',
                      type=str, default='./build/ecn.p4.p4info.txt')
    parser.add_argument('--bmv2-json', help='BMv2 JSON文件路径',
                      type=str, default='./build/ecn.json')
    args = parser.parse_args()
    
    if not os.path.exists(args.p4info):
        print(f"错误: 找不到p4info文件 {args.p4info}")
        sys.exit(1)
    if not os.path.exists(args.bmv2_json):
        print(f"错误: 找不到BMv2 JSON文件 {args.bmv2_json}")
        sys.exit(1)
    
    main(args.p4info, args.bmv2_json)